import os
import xml.etree.ElementTree as ET
from PIL import Image
from glob import glob

# 类别名到 YOLO class_id 的映射
CLASS_NAME_TO_ID = {
    'car': 0, 'truck': 1, 'pedestrian': 2, 'traffic_cone': 3, 'forklift': 4,
    'tractor': 5, 'flatbed': 6, 'trailer': 7, 'cyclist': 8, 'bicycle': 9,
    'water_horse': 10, 'guardrail': 11, 'unknown': 12
}

ROOT_DIR = 'datasets/scene_test'

# 遍历所有 scene-*image 子文件夹
scene_dirs = glob(os.path.join(ROOT_DIR, 'scene-*'))

for scene_dir in scene_dirs:
    if not os.path.isdir(scene_dir):
        continue

    print(f"\n🔍 正在处理：{scene_dir}")

    image_dir = os.path.join(scene_dir, 'images')
    label_dir = os.path.join(scene_dir, 'labels')
    annotation_file = os.path.join(scene_dir, 'annotations.xml')

    if not os.path.exists(annotation_file):
        print(f"[跳过] 无标注文件: {annotation_file}")
        continue

    os.makedirs(label_dir, exist_ok=True)

    # 解析 XML 文件
    tree = ET.parse(annotation_file)
    root = tree.getroot()

    for image_tag in root.findall('image'):
        file_name = os.path.basename(image_tag.get('name'))
        image_path = os.path.join(image_dir, file_name)

        if not os.path.exists(image_path):
            print(f"[跳过] 图像不存在: {file_name}")
            continue

        with Image.open(image_path) as img:
            img_width, img_height = img.size

        label_file = os.path.join(label_dir, os.path.splitext(file_name)[0] + '.txt')
        yolo_lines = []

        for box in image_tag.findall('box'):
            label = box.get('label')
            if label not in CLASS_NAME_TO_ID:
                print(f"[跳过] 未知类别: {label}")
                continue
            class_id = CLASS_NAME_TO_ID[label]

            xtl = float(box.get('xtl'))
            ytl = float(box.get('ytl'))
            xbr = float(box.get('xbr'))
            ybr = float(box.get('ybr'))

            x_center = (xtl + xbr) / 2 / img_width
            y_center = (ytl + ybr) / 2 / img_height
            width = (xbr - xtl) / img_width
            height = (ybr - ytl) / img_height

            yolo_line = f"{class_id} {x_center:.6f} {y_center:.6f} {width:.6f} {height:.6f}"
            yolo_lines.append(yolo_line)

        with open(label_file, 'w') as f:
            f.write('\n'.join(yolo_lines))

        print(f"[完成] {file_name} => {label_file}")
